# 安装依赖（需提前安装）
# pip install backtrader tushare matplotlib pandas

import tushare as ts
import backtrader as bt
import pandas as pd
import numpy as np
from datetime import datetime

# 设置Tushare Token（需自行申请）
ts.set_token('b91f8e85973449023ab471c313af892fed91df1fab7fc669abc3de8f')
pro = ts.pro_api()


# 修正后的EnhancedKDJ类
class EnhancedKDJ(bt.Indicator):
    """
    动态更新的KDJ指标（兼容最新版backtrader）
    """
    lines = ('K', 'D', 'J')
    params = (
        ('n', 9),  # RSV周期
        ('m', 3)  # 平滑周期
    )

    def __init__(self):
        # 计算n日最高最低价
        highest = bt.ind.Highest(self.data.high, period=self.p.n)
        lowest = bt.ind.Lowest(self.data.low, period=self.p.n)

        # RSV计算（包含除零保护）
        numerator = self.data.close - lowest
        denominator = highest - lowest
        rsv_raw = numerator / (denominator + 1e-7)  # 防止除零
        rsv = bt.If(denominator == 0, 50, 100 * rsv_raw)  # 处理极值情况

        # 正确的EMA计算方式（绑定到lines）
        self.lines.K = bt.ind.EMA(rsv, period=self.p.m)
        self.lines.D = bt.ind.EMA(self.lines.K, period=self.p.m)

        # J线直接使用公式计算
        self.lines.J = 3 * self.lines.K - 2 * self.lines.D

        # 显式声明依赖关系（可选，增强可读性）
        self.set_dependencies()

    def set_dependencies(self):
        """确保所有数据依赖正确建立"""
        self.addminperiod(self.p.n)


class KDJStrategyV2(bt.Strategy):
    params = (
        ('overbought', 83),  # 优化超买阈值[6](@ref)
        ('oversold', 17),  # 优化超卖阈值[6](@ref)
        ('risk_pct', 0.02),  # 单笔风险2%
        ('use_j', True)  # 启用J值过滤
    )

    def __init__(self):
        self.kdj = EnhancedKDJ()
        self.cross = bt.indicators.CrossOver(self.kdj.K, self.kdj.D)

        # 添加成交量过滤器[4](@ref)
        self.vol_ma20 = bt.indicators.SMA(self.data.volume, period=20)

    def next(self):
        # print(f"Date: {self.data.datetime.date(0)}, K={self.kdj.K[0]:.2f}, D={self.kdj.D[0]:.2f}")
        if not self.position:
            if self.cross > 0  :
                size = self.cal_position()
                self.buy(size=size)
        else:
            exit_cond = (self.cross < 0) or \
                        (self.kdj.K[0] > self.p.overbought) or \
                        (self.p.use_j and self.kdj.J[0] > 100)
            if exit_cond:
                self.sell(size=self.position.size)

    def cal_position(self):

        risk_amount = self.broker.getvalue() * self.p.risk_pct
        price = self.data.close[0]
        shares = int(risk_amount // price)
        return (shares // 100) * 100  # A股整手交易


def get_enhanced_data(ts_code, start, end):
    """数据获取与预处理（带异常值过滤）[3](@ref)"""
    df = pro.daily(ts_code=ts_code, start_date=start, end_date=end, adj='qfq')

    # 列名标准化
    df.rename(columns={
        'trade_date': 'date',
        'vol': 'volume'
    }, inplace=True)

    # 过滤异常K线
    df = df[(df['high'] >= df['low']) &
            (df['close'].pct_change().abs() < 0.2)]  # 剔除涨跌幅>20%的异常数据

    # 格式转换
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)
    return df.sort_index(ascending=True)[['open', 'high', 'low', 'close', 'volume']]


if __name__ == '__main__':
    # 参数配置
    symbol = '000063.SZ'  # 中兴通讯
    start_date = '20240101'
    end_date = '20250324'

    # 获取数据
    data = get_enhanced_data(symbol, start_date, end_date)

    # 初始化回测引擎
    cerebro = bt.Cerebro()

    # 添加数据
    data_feed = bt.feeds.PandasData(
        dataname=data,
        datetime=None,
        open=0, high=1, low=2, close=3, volume=4,
        openinterest=-1
    )
    cerebro.adddata(data_feed)

    # 添加策略
    cerebro.addstrategy(KDJStrategyV2)

    # 资金配置
    cerebro.broker.setcash(1000000)
    cerebro.broker.setcommission(
        commission=0.001,  # 0.1%佣金
        margin=1.0,
        automargin=False
    )

    # 添加分析器
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')

    # 执行回测
    results = cerebro.run()
    strat = results[0]



    # 可视化
    cerebro.plot(style='candlestick', barup='red', bardown='green')